Adds reduce_scatter into torchft
#102
Conversation
facebook-github-bot
added
the
CLA Signed
torchft/process_group_test.py
Outdated
| return True | ||
| return False | ||
| else: # cpu | ||
| if collective_str in ["reduce_scatter", "all_to_all"]: |
There was a problem hiding this comment.
oh wow -- didn't realize we don't support these on Gloo, good to know! cc @c-p-i-o
There was a problem hiding this comment.
ye, we miss many APIs on Gloo.
There was a problem hiding this comment.
this approach seems nice and explicit. but is it possible to instead just try: the test, and except: some specific NYI error? (i'm not sure if we raise a consistent type of NYI exception from backends?)
There was a problem hiding this comment.
this approach seems nice and explicit. but is it possible to instead just try: the test, and except: some specific NYI error? (i'm not sure if we raise a consistent type of NYI exception from backends?)
Yeah this is a good idea, I modified the block as follows:
for coll_str, args in collectives:
try:
coll = getattr(pg, coll_str)
work = coll(*args)
works[coll_str] = work
work.wait()
fut = work.get_future()
fut.wait()
# Check that all tensor arguments have the expected shapes and dtypes
check_tensors(args)
except RuntimeError as e:
if f"does not support {coll_str}" in str(e):
# Skip collectives that are not supported by the backend.
continue
raise e
|
Updated the test to be simpler, so I removed the utility functions I previously added. This should remove the need for a test refactor. I wanted to parameterize by collective, but #103 shows that tests got much slower after doing this. I will deprecate #103. I have also added an explicit |
What does this PR do?
Partially addresses #97 by adding
reduce_scatterintotorchft.Concretely, this consists of a few pieces:
reduce_scatterinto theProcessGroupfollowing the signature [here](https://github.com/pytorch/pytorch/blob/11f69808c64a65c68a4452250ba7719dcff27c78/torch/csrc/distributed/c10d/PyProcessGroup.hpp#L203ProcessGroup*we essentially follow the behavior of other collectives:ProcessGroupWrapper, it depends on the parent implementationProcessGroupDummy, it writes from the first input into outputProcessGroupBaby, it asserts inputs and moves underlying storage into shared memoryReduceScatterOptionsin_PickleSafeOptionsreduce_scatteras an option in_test_pg, however this necessitated a new function (named_should_run_collective) which was needed as e.g. GLOO does not supportreduce_scatter. This function essentially takes the collective, backend and device and copies the logic of the published supported collective matrix.Tests
Presubmits, and:
Next steps
The logic of
_should_run_collectiveis a bit confusing, as it allows "non defined backends" likeErrorSwallowing*through, to mimic the old behavior before this change. Testing here could become a bit unwieldy as we add more collectives and so a future step could be to refactor the testing.One nice change could be to parameterize tests by the collective. This will make potentially failing collectives more explicit and will reduce the time it takes to run individual tests. Likely can do this in the next PR.